import pdfplumber
import re


def academic_pdf_to_clean_txt(pdf_path, txt_path):
    """
    从学术论文PDF中提取文本，优化格式并保存为TXT文件
    包含空格优化、段落格式修复等处理
    """
    full_text = ""

    # 1. 提取PDF文本内容
    with pdfplumber.open(pdf_path) as pdf:
        for i, page in enumerate(pdf.pages):
            text = page.extract_text(
                x_tolerance=3,  # 水平间距容忍度
                y_tolerance=3,  # 垂直间距容忍度
                layout=True,  # 保持布局结构
                use_text_flow=True  # 保持文本流顺序
            )

            if text:
                # 应用学术文本优化
                text = optimize_academic_text(text, i)
                full_text += text + "\n\n"

    # 2. 应用空格优化处理
    cleaned_text = optimize_whitespace(full_text)

    # 3. 写入优化后的文本文件
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(cleaned_text)
    print(f"文本提取完成，已保存至: {txt_path}")


def optimize_academic_text(text, page_num):
    """
    优化学术论文文本格式
    """
    # 修复错误的换行（连接被断开的单词）
    text = re.sub(r"(\w)-\s+(\w)", r"\1\2", text)

    # 修复多栏文本中的错误换行
    if page_num > 0:  # 跳过标题页
        text = re.sub(r"(\w)\s*\n\s*(\w)", r"\1 \2", text)

    # 特殊格式处理
    text = re.sub(r"\$\{\}^{\text{(\w+)}}\$", r"^{\1}", text)  # 作者标记
    text = re.sub(r"(\d)\s+例", r"\1例", text)  # 数字连接
    text = re.sub(r"P\s*<\s*0\.0\s*5", r"P<0.05", text)  # 统计标记

    # 处理表格中的特殊符号
    text = re.sub(r"(\d)\s*±\s*(\d)", r"\1±\2", text)

    # 保留标题格式
    if page_num == 0:
        text = re.sub(r"·综合研究·", "\n·综合研究·\n", text)
        text = re.sub(r"摘要\s*目的", "摘要\n目的", text)

    # 处理参考文献格式
    if "参考文献" in text:
        text = re.sub(r"($$\d+$$)\s*([A-Z])", r"\1 \2", text)

    return text


def optimize_whitespace(text):
    """
    优化文本中的空格问题
    """
    # 替换所有制表符为空格
    text = text.replace('\t', ' ')

    # 删除每行首尾的空格（包括连续空格）
    text = re.sub(r'^[ \t]+|[ \t]+$', '', text, flags=re.MULTILINE)

    # 将连续两个及以上的空格替换为一个空格[1](@ref)
    text = re.sub(r' {2,}', ' ', text)

    # 保留必要的学术格式空格
    text = re.sub(r'([a-zA-Z])([\u4e00-\u9fa5])', r'\1 \2', text)  # 英文后接中文加空格
    text = re.sub(r'([\u4e00-\u9fa5])([a-zA-Z])', r'\1 \2', text)  # 中文后接英文加空格
    text = re.sub(r'(\d)\s*±\s*(\d)', r'\1 ± \2', text)  # 保留统计符号空格

    # 修复被错误分割的段落
    text = re.sub(r'([。！？；："”])\s*\n\s*', r'\1\n\n', text)

    return text


# 使用示例
if __name__ == "__main__":
    input_pdf = "trends in incidence of hip fracture source jama netw open so 2024 may 1 7 5 e249186 的全文翻译 - 副本.pdf"
    output_txt = "trends in incidence of hip fracture source jama netw open so 2024 may 1 7 5 e249186 的全文翻译 - 副本.txt"
    academic_pdf_to_clean_txt(input_pdf, output_txt)
